from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from datasets import load_dataset
from peft import LoraConfig, get_peft_model
import torch

# ====== 第1步：加载本地模型路径 ======
#model_dir = "E:/models/qwen/qwen--Qwen3-0.6B"  # 这里改成你下载路径
model_dir = "D:\models\qwen\qwen\Qwen3-0___6B"

# ====== 第2步：加载分词器和模型 ======
tokenizer = AutoTokenizer.from_pretrained(model_dir, use_fast=False)  # use_fast=False防止兼容问题
# 加载模型，自动放到显卡上（自动判断CUDA）
model = AutoModelForCausalLM.from_pretrained(
    model_dir,
    device_map="auto",          # 自动分配设备
    dtype=torch.float16   # 使用半精度节省显存
)

# ====== 第3步：启用LoRA低显存微调 ======
# LoRA是Low-Rank Adaptation，轻量化微调方式，不改动原权重
lora_config = LoraConfig(
    r=8,                # 秩，越大可学习能力越强，但显存占用略增
    lora_alpha=16,      # 缩放系数
    target_modules=["q_proj", "v_proj"],  # 只在注意力层做适配
    lora_dropout=0.05,  # 防止过拟合
    bias="none",        # 不训练偏置项
    task_type="CAUSAL_LM"  # 因果语言建模任务（对话生成类）
)
model = get_peft_model(model, lora_config)

# ====== 第4步：加载数据集 ======
# data.jsonl 每行是 {"instruction": "...", "output": "..."}
dataset = load_dataset("json", data_files="../data/data_1.jsonl", split="train")

# ====== 第5步：将每条数据拼接成模型能读的格式 ======
def tokenize_fn(example):
    # 拼接指令与回答
    prompt = f"用户：{example['instruction']}\n小易哥：{example['output']}"
    # 分词，截断防止太长
    tokens = tokenizer(prompt, truncation=True, max_length=512)
    # labels 就是 input_ids 的副本，用于训练目标
    tokens["labels"] = tokens["input_ids"].copy()
    return tokens

tokenized_dataset = dataset.map(tokenize_fn)

# ====== 第6步：设置训练参数 ======
training_args = TrainingArguments(
    output_dir="./qwen-finetune-output",  # 输出目录
    per_device_train_batch_size=1,        # 每张显卡的batch大小（6G显存只能1）
    gradient_accumulation_steps=4,        # 累积梯度，相当于逻辑batch=4
    num_train_epochs=3,                   # 训练轮数
    learning_rate=1e-4,                   # 学习率
    fp16=True,                            # 开启半精度训练
    logging_steps=1,                      # 每步打印日志
    save_strategy="epoch",                # 每轮保存一次模型
    save_total_limit=2,                   # 最多保留2个checkpoint
    report_to="none"                      # 不使用wandb
)

# ====== 第7步：创建Trainer对象 ======
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset
)

# ====== 第8步：启动训练 ======
trainer.train()
